import logging, sys
import time
import grpc
import os
import numpy as np
from concurrent import futures


# from proto import detect_pb2
# from proto import detect_pb2_grpc
import luoding_src.proto.detect_pb2 as detect_pb2
import luoding_src.proto.detect_pb2_grpc as detect_pb2_grpc
import cv2
from ultralytics import YOLO

logging.basicConfig(level=logging.INFO,  # 打印的日志级别
                    format='%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s'  # 日志格式
                    )

def mat2bytes(image: np.ndarray) -> bytes:
    return np.array(cv2.imencode('.jpg', image)[1]).tobytes()

def bytes2mat(image: bytes) -> np.ndarray:
    return cv2.imdecode(np.array(bytearray(image), dtype='uint8'), cv2.IMREAD_UNCHANGED)

class DetectServer(detect_pb2_grpc.DetectServiceServicer):

    def __init__(self):
        self._logger = logging.getLogger('model loading...')
        # self._logger.setLevel(logging.INFO)
        self.model = YOLO('last.pt')
        self.img_default = cv2.imread('default.jpg')
        # self.img_default = cv2.imread('1.jpg')
        self.img = self.img_default
        # self.img_bytes = self.img.tobytes()
        print('model load completed')

    def DetectOperation(self, request, context):
        processed_img = self.img_default
        processed_img = cv2.resize(processed_img, (640, 400))
        abnormal_signal = True
        results = 2
        logging.info('one request...')
        # print(processed_img.shape, abnormal_signal, results)

        try:
            # todo 空req
            if not request:
                self._logger.info('request is None')
                return detect_pb2.DetectResponse(processed_img=mat2bytes(processed_img),
                                                abnormal_signal=abnormal_signal,
                                                results=results)

                # todo 图片处理
            processed_img = self.img
            result_yolo = self.model(self.img)
            try:
                for r in result_yolo:
                    processed_img = r.plot()
                    cls = int(r.boxes.cls.numpy()[0])
                    print(cls)
                    if cls == 0:
                        abnormal_signal = False
                        results = 0
                    if cls == 1:
                        abnormal_signal = True
                        results = 1
            except:
                    abnormal_signal = True
                    results = 2


            return detect_pb2.DetectResponse(processed_img=mat2bytes(processed_img), abnormal_signal=abnormal_signal,
                                                results=results)
        except:

            self._logger.info('DetectOperation Error!')
            return detect_pb2.DetectResponse(processed_img=mat2bytes(processed_img), abnormal_signal=abnormal_signal,
                                            results=results)

    def streamVideo(self, request, context):
        try:
            self.img_bytes = request.origin_frame
            self.img = bytes2mat(self.img_bytes)
            # logging.info('accept one frame!')
            return detect_pb2.CommonResponse(code=True)
        except:
            return detect_pb2.DetectResponse(code=False)


def run():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    detect_pb2_grpc.add_DetectServiceServicer_to_server(DetectServer(), server)
    logging.info('start service...')
    print('start service...')
    server.add_insecure_port('[::]:6667')
    server.start()
    logging.info('service running...')
    print('service running...')
    try:
        while True:
            time.sleep(60 * 60 * 24)
    except KeyboardInterrupt:
        server.stop(0)
        logging.info('stop service...')


if __name__ == '__main__':
    run()
